import h5py
import numpy as np
import matplotlib.pyplot as plt
from mpl_toolkits.mplot3d import Axes3D
from sklearn.cluster import DBSCAN
import math

#INPUT_DATA_DIR = '/home/behollis/DATA/pierre/ocean/'
#OUTPUT_DATA_DIR = '/home/behollis/Dropbox/visweek2015/oslines/'

OUTPUT_DATA_DIR = '/home/brad/visweek2015-revision-code/oslines/'
INPUT_DATA_DIR = '/home/brad/DATA/ts00050/'
SLINESF = 'lockSlines.1100steps.hdf5'

TOTMEMS = 30
X = 125
Y = 125
DIM = 2

'''
implementation of entropy based streamline cluster from:

Chen, ChengKai, et al. An illustrative visualization framework for 
3d vector fields. Computer Graphics Forum. 
Vol. 30. No. 7. Blackwell Publishing Ltd, 2011.
'''

import numpy.linalg as la

def writeElapsedTime(start, end):
    import os
    systime = os.times()
    tout = open(OUTPUT_DATA_DIR+'timinglockSlineClusters.10pts.{0}.txt'.format(systime[4]), 'w') #sys time
    tout.write('Elapsed time: {0}'.format(end - start))
    tout.close() 
    
    
from numpy import (array, dot, arccos, clip)
from numpy.linalg import norm
 
def py_ang(v1, v2):
    ''' Returns the angle in radians between vectors 'v1' and 'v2' '''
    
    if norm(v1) == 0 or norm(v2) == 0:
        return 0.0
    
    cosang = np.dot(v1, v2) / norm(v1) / norm(v2) 
    angle = arccos(clip(cosang, -1, 1))
    
    return  angle

def calcSlineEntropies(slines):
    ''' returns avg linear and angular entropies for a seed '''
    
    aentropies = list()
    lentropies = list()
    
    for sl in slines: 
        
        # trim slines back to 500 points forward / backward
        sl[0] = sl[0][500:-500]
        sl[1] = sl[1][500:-500]
            
        total_len, sl_seg_lengths = calcSlineLength(sl) 
        
        #number of line segments
        m = len( sl_seg_lengths ) 
       
        #print 'sline lengths: ' + str(total_len)
        
        TOL_MIN = 0.000000001 
        #summation of linear entropies
        sum_lentropy = 0.0
        for s in sl_seg_lengths:
            #print s
            if s > TOL_MIN:
                sum_lentropy += ( s / total_len ) * math.log(s / total_len, 2)  
       
        lentropy = -1.0 * sum_lentropy / math.log( m + 1,2 )
       
        #find total angular variation and angle between each segment
        sl_seg_ang = list()
        total_ang = 0.0
        #total_ang_change = 0.0
        num_pts = len(sl[0])
        for j in range( 1, num_pts - 1 ):
            pt0 = ( sl[0][j-1], sl[1][j-1], 0.0 )
            pt1 = ( sl[0][j]  , sl[1][j]  , 0.0 )
            pt2 = ( sl[0][j+1], sl[1][j+1], 0.0 )
           
            v0 = np.array(pt1) - np.array(pt0)
            v1 = np.array(pt2) - np.array(pt1)
           
            ang = py_ang(v0,v1)
            #print 'ang: ' + str(ang)
           
            sl_seg_ang.append(ang)
           
            total_ang += math.fabs(ang)
            #total_ang_change += ang
           
        #summation of linear entropies
        sum_aentropy = 0.0
        for a in sl_seg_ang:
            if a != 0.0:
                sum_aentropy += ( a / total_ang ) * math.log(a / total_ang, 2) 
               
        aentropy = -1.0 * sum_aentropy / math.log( m, 2 )
       
        aentropies.append(aentropy)
        lentropies.append(lentropy)
       
    numslines = len(slines)
    sumAE = np.sum(np.asarray(aentropies))
    sumLE = np.sum(np.asarray(lentropies))
    
    return sumLE / numslines, sumAE / numslines
       
def calcSlineLength(sl):
    ''' find total length of streamline and length of each seg '''
    sl_seg_lengths = list()
    total_len = 0.0
    num_pts = len(sl[0])
    for j in range( 0, num_pts - 1): #trim extra advection off
        pt0 = ( sl[0][j], sl[1][j], 0.0 )
        pt1 = ( sl[0][j+1], sl[1][j+1], 0.0 )
        
        dist = np.linalg.norm( np.array(pt1) - np.array(pt0) )
        #print 'dist: ' + str(dist)
        sl_seg_lengths.append( dist )
        
        total_len += dist
        
    return total_len, sl_seg_lengths
       
def calcStreamlineFeatures(slines, pts_to_sample=0.0):
    slines_features = list()
    
    for sl in slines: 
        
        # trim slines back to 500 points forward / backward
        sl[0] = sl[0][500:-500]
        sl[1] = sl[1][500:-500]
        
        if len(sl[0]) < 3 or len(sl[1]) < 3:
            continue 
        
         
        slf = list() #list of streamline feature vectors to cluster
        #find first, middle, and last points on sl
        begin_pt = [ sl[0][0], sl[1][0] ] 
        half_num_pts = len(sl[0])/2
           
        end_pt = [ sl[0][-1], sl[1][-1] ] 
        slf = begin_pt + end_pt 
           
        linear_entropy = 0.0
        angl_entropy = 0.0
           
        num_pts = len( sl[0] ) 
        num_segs = num_pts - 1
           
        total_len, sl_seg_lengths = calcSlineLength(sl)
               
        #find midpoint location of arc length
        mid_pt_idx = 0
        arc_length = 0.0
        
        for j in range( 0, num_pts - 1 ):
            pt0 = ( sl[0][j], sl[1][j], 0.0 )
            pt1 = ( sl[0][j+1], sl[1][j+1], 0.0 )
           
            dist = np.linalg.norm( np.array(pt1) - np.array(pt0) )
            arc_length += dist
          
            if arc_length >= ( total_len / 2.0 ):
                mid_pt_idx = j
                break
           
        # add mid point
        mid_pt = [ sl[ 0 ][ mid_pt_idx ], sl[ 1 ][ mid_pt_idx ] ] 
        slf += mid_pt
       
        # find additional points to be sampled
        arc_length = 0.0
        pt_count = 0.0
        if pts_to_sample > 0:
            for j in range( 0, num_pts - 1 ):
                pt0 = ( sl[0][j], sl[1][j], 0.0 )
                pt1 = ( sl[0][j+1], sl[1][j+1], 0.0 )
               
                dist = np.linalg.norm( np.array(pt1) - np.array(pt0) )
                arc_length += dist
              
                if arc_length >= pt_count * ( total_len / pts_to_sample ):
                    if pt_count >= pts_to_sample:
                        break
                   
                    pt_count += 1
                    curr_pt = [ sl[0][j], sl[1][j] ] 
                    slf += curr_pt 
                  
        slines_features.append( slf )
    
    # return feature vectors of streamlines
    return slines_features
        
def readStreamlines(x, y, f):
    slines = list()
    groups = f.keys()
    
    for mem in groups:
        dir = str(mem).zfill(4) + '/x' + str(x).zfill(3) + '/y' + str(y).zfill(3)
    
        try:
            xlst = list(f[dir][0])
            ylst = list(f[dir][1])
            zlst = 0.0 * len(list(f[dir][0]))
            
            slines.append([xlst,ylst,zlst])
        except:
            continue
    
    return slines
        
if __name__ == '__main__':
    
    import time
    start = time.time()
    
    f = h5py.File(OUTPUT_DATA_DIR+SLINESF, 'r')
    
    output = h5py.File( OUTPUT_DATA_DIR+'lockEntropyONLY.fixed.hdf5', 'w')
    #dslclus = output.create_dataset(name='slineclusters', shape=(X,Y), dtype='f')
    #dsampled = output.create_dataset(name='ptsSampledForClustering', shape=(X,Y), dtype='f')
    dentroL = output.create_dataset(name='avgAngularEntropy', shape=(X,Y), dtype='f')
    dentroA = output.create_dataset(name='avgLinearEntropy', shape=(X,Y), dtype='f')
    
    PTS_SAMPLED_BEYOND_MIN = 10
    
    FEATURES = 2 * 3 + 2 * PTS_SAMPLED_BEYOND_MIN #default of first, middle and end point in 2 dimensions
   
    for x in range(0,X):
        for y in range( 0, Y):
            print 'calculating {0}, {1}'.format(x, y)

            sltotal = readStreamlines(x,y,f)
         
            avg_lin_entropy = 0.0; avg_ang_entropy = 0.0
            try:
                avg_lin_entropy, avg_ang_entropy = calcSlineEntropies(sltotal)
                if math.isnan(avg_lin_entropy) :
                    avg_lin_entropy = 0.0
                
                if math.isnan(avg_ang_entropy) :
                    avg_ang_entropy = 0.0 
            except:
                print 'error in calculating entropies...'
             
            dentroL[x,y] = avg_lin_entropy
            dentroA[x,y] = avg_ang_entropy
            
            print 'ang entropy: ' + str(avg_ang_entropy)
            print 'lin_entropy: ' + str(avg_lin_entropy)
             
        
            '''
            #calculate feature vectors for streamlines
            try:
                feat = calcStreamlineFeatures(sltotal, pts_to_sample=PTS_SAMPLED_BEYOND_MIN)
            
                feat_total = np.array( feat )
                feat_total.reshape(len(feat_total), FEATURES)
                
                # STIRRING PARAMS
                #dcov[x,y] = np.cov(np.array(tpts).T)#pca.get_covariance()
                # let's use eps to be some signficant fraction of
                # the simulation physical domain...we have 152 x 152 cells
                # five percent of the diagonal distance across the domain
                # thus, diag dist = 176.67
                # 0.05 * 214.96 = 8.38
                # NOTE: include this metric and value in the paper
                # min samples is 0.1 of total members, or 2 
                
                #db = DBSCAN(eps=8.38, min_samples = 2).fit( np.asarray(feat_total) )
                db = DBSCAN(eps=8.38, min_samples = 2).fit( np.asarray(feat_total) )

                
                core_samples_mask = np.zeros_like(db.labels_, dtype=bool)
                core_samples_mask[ db.core_sample_indices_] = True
                labels = db.labels_
        
                # Number of clusters in labels, ignoring noise if present.
                n_clusters_ = len(set(labels)) - (1 if -1 in labels else 0)
                dslclus[x,y] = n_clusters_
                #print n_clusters_
            except:
                print 'exception occured when trying to find clusters!'
                dslclus[x,y] = 0
                
            dsampled[x,y] = FEATURES
            print 'clusters: ' + str(dslclus[x,y])
            '''
            
    end = time.time()
    writeElapsedTime(start, end)

    f.close()
    output.close()
    
    print 'finished!'